/*
 * anls_sequential.cc
 *
 * Copyright 2013 Long Tran (ltran3@gatech.edu)
 *
 * Sequential ALS, solving saddle point of the Lagrangian
 *   L = sum |A - WH^T|^2 + <Lambda, W-U> + <Pi, H-V>
 *        + 0.5*rho*(|W-U|^2+|H-V|^2)
 *   subject to U >= 0, V >= 0
 * 
 * This program is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation; either version 2 of the License, or
 * (at your option) any later version.
 * 
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 * 
 * You should have received a copy of the GNU General Public License
 * along with this program; if not, write to the Free Software
 * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston,
 * MA 02110-1301, USA.
 * 
 */

 #include <armadillo>

#include "helper/helper.h"
#include "algorithm/StopWatch.h"
#include "data/FileAccess.h"
#include <boost/mpi.hpp>
extern boost::mpi::communicator* p_world;
extern po::variables_map vm;
#include "mpi_helper.h"
#include "arma_helper.h"

using namespace arma;

class SequentialANLS
{
public:
  mat A, W, H, H_old;
  std::vector<Iteration> iterInfo;
  
  double rho;
  int m, n, k;
  mat U, V, U_old, V_old, Lambda, Pi, Delta;

  int maxiter, interval;
  bool normalizing;
  double rho_incr, rho_decr;

  StopWatch watch, clock;
  TimeDuration runtime;    
public:
  SequentialANLS()
  {
    PO_REQUIRE(vm, "k");
    PO_VARIABLE_VALUE(rho, vm, "rho", double); // ADMM parameter
    PO_VARIABLE_VALUE(rho_incr, vm, "rho-incr", double); // ADMM parameter
    PO_VARIABLE_VALUE(rho_decr, vm, "rho-decr", double); // ADMM parameter
    PO_VARIABLE_VALUE(maxiter, vm, "iter", int);
    PO_VARIABLE_VALUE(interval, vm, "interval", int);
    PO_VARIABLE_VALUE(k, vm, "k", int);
  }

  void LoadData()
  {
    PO_REQUIRE(vm, "data");
    PO_VARIABLE(data, vm, "data", std::string);
    
    std::stringstream ss;
    ss << data << "_" << MPI_RANK;
    A.load(ss.str(), arma_binary);
    m = A.n_rows; 
    n = A.n_cols;
  }

  void UpdateWH()
  {
    mat HTHpI = trans(H)*H+rho*eye(k,k);
    W = trans(solve(HTHpI, trans(A*H-Lambda+rho*U)));
    
    // update H
    H_old = H;
    mat WTWpI = trans(W)*W+rho*eye(k,k);
    H = trans(solve(WTWpI, trans(trans(A)*W-Pi+rho*V)));
  }
  
  void UpdateUV()
  {
    U_old = U;
    U = W + Lambda/rho;
    zero_threshold(U);

    V_old = V;
    V = H + Pi/rho;
    zero_threshold(V);
  }

  void UpdateLambdaPi()
  {
    Lambda += rho*(W-U);
    Pi += rho*(H-V);
  }

  void UpdateRho() 
  {
    // primal feasibility
    double w_norm = norm(W-U, "fro"), h_norm = norm(H-V, "fro");
    double primal = sqrt(w_norm*w_norm + h_norm*h_norm);
    
    // dual feasibility
    double v_norm = norm(rho*(V-V_old), "fro");
    mat U_NORM = /*A*(H_old-H) 
		   - W*( trans(H_old)*H_old - trans(H)*H ) */
      - rho * (U_old - U);
    double u_norm = norm( U_NORM, "fro" ); 
    double dual = sqrt(u_norm*u_norm + v_norm*v_norm);

    if (primal > 10*dual) rho *= rho_incr;
    else if (dual > 10*primal) rho *= rho_decr;
    else rho = rho;
    if (rho > 1e3) rho = 1e3;
    if (rho < 1e-2) rho = 1e-2;
  }

  void RecordProgress(int iter)
  {
    double res = norm(A-W*trans(H),"fro");
    Iteration info(iter, res, runtime, runtime);
    iterInfo.push_back(info);

    if (iter % interval == 0) {
      MSG("it=" << iter 
	  << " |res|=" << res
	  << " |W|=" << norm(U, "fro")
	  << " |H|=" << norm(V, "fro")
	  << " d=" << SECONDS(info.duration)
	  << " |fea|=" << norm(W-U,"fro") + norm(H-V, "fro")
	  << " rho=" << rho
	  );
    }
  }
  
  void SolveANLS()
  {
    watch.Reset();
    clock.Reset();

    W.randu(m, k);
    H.randu(n, k);
    U = zeros<mat>(m, k);
    V = zeros<mat>(n, k);
    Lambda = zeros<mat>(m, k);
    Pi = zeros<mat>(n, k);

    for (int iter = 0; iter < maxiter; iter++) {
      RecordProgress(iter);
      // update W, H
      clock.GetDuration();
      UpdateWH();
      UpdateUV();
      UpdateLambdaPi();
      runtime += clock.GetDuration();

      UpdateRho();      
    }
    RecordProgress(maxiter);
    MSG(watch.GetTotalDuration());
    W = U;
    H = V;
  }
};

void run_anls()
{
  MSG("Sequential ANLS");
  SequentialANLS a;
  a.LoadData();
  a.SolveANLS();

  if (vm.count("iterinfo")) {
    PO_VARIABLE(infoFile, vm, "iterinfo", std::string);
    FileAccessWrite(out, infoFile, NULL);
    const std::vector<Iteration>& info = a.iterInfo;
    for (unsigned int i = 0; i < info.size(); i++) {
      out << info[i].iter << "\t" << (double)info[i].duration.total_milliseconds()/1000 << "\t" 
	  << (double)info[i].compute.total_milliseconds()/1000 << "\t" 
	  << (double)info[i].communication.total_milliseconds()/1000 << "\t"
	  << info[i].residual << std::endl;
    }
  }
}

/*
// sequential ALS algorithm
void solve_anls(const arma::mat& A, arma::mat& W, arma::mat& H, std::vector<Iteration>& iterInfo)
{
  StopWatch watch;
  using namespace arma;
  int m = A.n_rows, n = A.n_cols, k = W.n_cols;
  mat U = zeros<mat>(m, k), V = zeros<mat>(n, k), X = zeros<mat>(n, k), V_tmp = zeros<mat>(n, k);
  mat Lambda = zeros<mat>(m, k), Pi = zeros<mat>(n, k), Delta = zeros<mat>(n, k);

  PO_REQUIRE(vm, "iter");

  PO_VARIABLE(rho, vm, "rho", double); // ADMM parameter
  //PO_VARIABLE(alpha, vm, "rate", double); // dual ascent step
  PO_VARIABLE(maxiter, vm, "iter", int);
  PO_VARIABLE(interval, vm, "interval", int);

  StopWatch clock;
  TimeDuration runtime;
  for (int iter = 0; iter < maxiter; iter++) {
    // update W
    clock.GetDuration();
    mat HTHpI = trans(H)*H+rho*eye(k,k);
    W = trans(solve(HTHpI, trans(A*H-Lambda+rho*U)));

    // update H
    mat WTWpI = trans(W)*W+rho*eye(k,k);
    H = trans(solve(WTWpI, trans(trans(A)*W-Pi-Delta+rho*V+rho*X)));

    // update U, V, X
    U = W + Lambda/rho;
    V = H + Pi/rho;
    zero_threshold(U);
    zero_threshold(V);
    //if (vm.count("normalize")) normalize_column(V_tmp, V);
    if (vm.count("normalize")) {
      X = H + Delta/rho;
      normalize_column(X);
    }

    // update Lambda, Pi, Delta (dual variables)
    double rate = rho ; // alpha/sqrt(iter+1)*rho;
    Lambda += rate*(W-U);
    Pi += rate*(H-V);
    if (vm.count("normalize")) Delta += rate*(H-X);
    runtime += clock.GetDuration();

    if (iter % interval == 0 || iter == maxiter-1) {
      double res = norm(A-U*trans(V),"fro");
      MSG("it=" << iter 
	  << " |res|=" << res
	  << " |W|=" << norm(U, "fro")
	  << " |H|=" << norm(V, "fro")
          << " d=" << watch.GetDuration());
      Iteration info(iter, res);
      info.duration = clock.GetTotalDuration();
      info.compute = runtime;
      iterInfo.push_back(info);
    }
  }
  MSG(watch.GetTotalDuration());
  W = U;
  H = V;
}

void run_anls()
{
  PO_REQUIRE(vm, "data");
  PO_REQUIRE(vm, "k");
  PO_VARIABLE(data, vm, "data", std::string);
  PO_VARIABLE(k, vm, "k", int);

  using namespace arma;
  mat A;
  std::stringstream ss;
  ss << data << "_" << MPI_RANK;
  A.load(ss.str(), arma_binary);
  int m = A.n_rows, n = A.n_cols;

  srand(100);
  mat W0, H0;
  W0.randu(m, k);
  H0.randu(n, k);
  std::vector<Iteration> iterInfo;
  solve_anls(A, W0, H0, iterInfo);

  //MSG_VERBOSE(vm, "W0=\n" << W0 << "H0=\n" << H0);
  MSG("|res|=" << norm(A-W0*trans(H0), "fro"));
  MSG("sqrt(k)=" << sqrt(k) << " norm(H,fro)=" << norm(H0, "fro"));

  if (vm.count("iterinfo")) {
    PO_VARIABLE(infoFile, vm, "iterinfo", std::string);
    FileAccessWrite(out, infoFile, NULL);
    const std::vector<Iteration>& info =  iterInfo;
    for (unsigned int i = 0; i < info.size(); i++) {
      out << info[i].iter << "\t" << (double)info[i].duration.total_milliseconds()/1000 << "\t" 
	  << (double)info[i].compute.total_milliseconds()/1000 << "\t" 
	  << (double)info[i].communication.total_milliseconds()/1000 << "\t"
	  << info[i].residual << std::endl;
    }
  }
}
*/
